import os
import cv2
import numpy as np
from pathlib import Path
from tqdm import tqdm
from color import *

PROGRESS_BAR = False


def smart_print(*a, **b):
    if not PROGRESS_BAR:
        print(*a, **b)


def cv_imread(img_path, mode=1):
    """
    解决 OpenCV 无法载入中文图片的问题
    :param img_path: 图片路径
    :param mode: 模式，1为彩色，0为黑白
    :return:
    """
    img = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), mode)
    return img


def cv_imwrite(image: np.array, path: str, ext_name='png'):
    """
    解决 OpenCV 无法保存中文图片的问题
    :param image: 待保存图片
    :param path: 保存路径
    """
    cv2.imencode('.' + ext_name, image)[1].tofile(path)


def crop_to_square_at_center(image: np.array):
    """
    将图片裁剪成正方形
    :param image: 图片
    :return: 裁剪后的图片
    """
    height, width = image.shape[:2]
    if height == width:
        return image
    elif height > width:
        return image[int((height - width) / 2):int((height + width) / 2), :]
    else:
        return image[:, int((width - height) / 2):int((width + height) / 2)]


def resize_square_img(img: np.array, width: int):
    img = crop_to_square_at_center(img)
    return cv2.resize(img, (width, width))


def add_alpha_channel(img):
    """ 为jpg图像添加alpha通道 """
    if img.ndim == 3:
        b_channel, g_channel, r_channel = cv2.split(img)  # 剥离jpg图像通道
        alpha_channel = np.ones(
                b_channel.shape, dtype=b_channel.dtype) * 255  # 创建Alpha通道

        img_new = cv2.merge(
                (b_channel, g_channel, r_channel, alpha_channel))  # 融合通道
        return img_new
    else:
        return img


def load_png_image(image_path: str):
    """
    加载png图片
    :param image_path: 图片路径
    :return:
    """
    img = cv_imread(image_path, cv2.IMREAD_UNCHANGED)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    return img


def add_one_to_another(bg: str, fg: str, save: str):
    """
    Copy the png file `fg` onto the background image `bg`,
    auto process size, format, etc...
    then save to `save`.
    """
    smart_print(f'Reading {G(bg)}')
    bgimg = cv_imread(bg, cv2.IMREAD_UNCHANGED)
    fgimg = cv_imread(fg, cv2.IMREAD_UNCHANGED)
    smart_print(f'Dimensions:', bgimg.ndim, fgimg.ndim)

    if bgimg.ndim == 2:
        bgimg = cv2.cvtColor(bgimg, cv2.COLOR_GRAY2BGR)
        bgimg = add_alpha_channel(bgimg)

    if fgimg.ndim != 3:
        smart_print(R('fgimg.ndim != 3! template must be PNG format.'))
        return

    smart_print(f'Resizing {bg}')
    bgimg = resize_square_img(bgimg, 400)
    fgimg = resize_square_img(fgimg, 400)

    # cv2.imwrite('bg1.png', bgimg)
    # cv2.imwrite('fg1.png', fgimg)
    # copy transparent fgimg onto bgimg

    # 获取要覆盖图像的alpha值，将像素值除以255，使值保持在0-1之间
    alpha_png = fgimg[:, :, 3] / 255.0
    alpha_jpg = 1 - alpha_png

    smart_print(f'Merging {bg}')
    for c in range(0, 3):
        bgimg[:, :, c] = ((alpha_jpg * bgimg[:, :, c]) +
                          (alpha_png * fgimg[:, :, c]))

    cv_imwrite(bgimg, save, 'png')
    smart_print(f'Done with {G(bg)}, save to {G(save)}\n')


def default_batch(bg_folder, template_path, output_folder):
    """
    Batch, default mode, use only one template.
    :param bg_folder:
    :param template_path:
    :param output_folder:
    :return:
    """
    global PROGRESS_BAR
    PROGRESS_BAR = True
    for i in tqdm(os.listdir(bg_folder)):
        file = Path(bg_folder) / i
        save = Path(output_folder) / Path(i.split('.')[0] + '_processed.png')
        if save.exists():
            continue
        add_one_to_another(str(file), template_path, str(save))


def multi_template_batch(bg_folder, template_folder, output_folder):
    """
    Batch, multi mode, use multiple templates, and apply these to every single file.
    :param bg_folder:
    :param template_folder:
    :param output_folder:
    :return:
    """
    global PROGRESS_BAR
    PROGRESS_BAR = True
    templates = os.listdir(template_folder)
    backgrounds = os.listdir(bg_folder)
    with tqdm(total=len(templates) * len(backgrounds)) as bar:
        for template in templates:
            for background in backgrounds:
                ifile = Path(bg_folder) / background
                ofile = Path(output_folder) / Path(background.split('.')[0] + '_with_' + template.split('.')[0] + '.png')
                if ofile.exists():
                    continue
                add_one_to_another(str(ifile), str(Path(template_folder) / template), str(ofile))
                bar.update(1)


if __name__ == '__main__':
    # add_one_to_another('test.png', 'cover.png', 'output.png')
    multi_template_batch('backgrounds', 'template', 'results')
